import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from .config import *
import os

def evaluate_model(model, test_data, times_test):
    """评估模型性能"""
    if config.MODEL_SELECT == "LSTM" or config.MODEL_SELECT == "CNN_LSTM" or config.MODEL_SELECT == "U-Net":
        images, pv_log, pv_pred= test_data
    elif config.MODEL_SELECT == "CLASSIFY_MODEL":
        images, pv_log, pv_pred, weather_data = test_data
    else:
        pass
    if config.MODEL_SELECT == "CNN_LSTM" and config.WEATHER_SELECT == "CLOUDY":
        loss = model.evaluate([images, pv_log], pv_pred, verbose=0)
        predictions = np.squeeze(model.predict(([images, pv_log]), batch_size=200, verbose=1))
    elif config.MODEL_SELECT == "CNN_LSTM" and config.WEATHER_SELECT == "SUNNY":
        loss = model.evaluate(pv_log, pv_pred, verbose=0)
        predictions = np.squeeze(model.predict((pv_log), batch_size=200, verbose=1))
    elif config.MODEL_SELECT == "CNN_LSTM" and config.WEATHER_SELECT == "OVERCAST":
        loss = model.evaluate(pv_log, pv_pred, verbose=0)
        predictions = np.squeeze(model.predict((pv_log), batch_size=200, verbose=1))
    elif config.MODEL_SELECT == "CNN_LSTM" and config.WEATHER_SELECT == "ALL":
        loss = model.evaluate(pv_log, pv_pred, verbose=0)
        predictions = np.squeeze(model.predict((pv_log), batch_size=200, verbose=1))
    elif config.MODEL_SELECT == "LSTM" or config.MODEL_SELECT == "MLP":
        loss = model.evaluate(pv_log, pv_pred, verbose=0)
        predictions = np.squeeze(model.predict((pv_log), batch_size=200, verbose=1))
    elif config.MODEL_SELECT == "CLASSIFY_MODEL":
        loss = model.evaluate([pv_log, weather_data], pv_pred, verbose=0)
        predictions = np.squeeze(model.predict(([pv_log, weather_data]), batch_size=200, verbose=1))
    rmse = np.sqrt(np.mean((predictions - pv_pred)**2))
    
    metrics = {
        'overall': calculate_metrics(predictions, pv_pred)
    }
    
    return metrics, predictions

def calculate_metrics(pred, true):
    mse = np.mean((pred - true)**2)
    mae = np.mean(np.abs(pred - true))
    return {'RMSE': np.sqrt(mse), 'MAE': mae, 'MSE': mse}

def plot_forecast_segments(weather_select, times, predictions, test_pv_pred, 
                           save_dir='./data/model_output/SUNSET_forecast_2017_2019_data/', 
                           max_time_gap=np.timedelta64(6, 'h')):
    # 创建保存目录（如果不存在）
    os.makedirs(save_dir, exist_ok=True)
    # 定义保存的文件名
    filename = f'{weather_select}_forecast_plot.png'
    save_path = os.path.join(save_dir, filename)
    # 找到时间间隔过大的位置
    time_diffs = np.diff(times)  # 计算相邻时间点之间的差值
    gap_indices = np.where(time_diffs > max_time_gap)[0]  # 找到时间间隔大于阈值的位置
    # 将数据分割成多个时间段
    segments = []
    start_idx = 0
    for idx in gap_indices:
        segments.append((times[start_idx:idx+1], predictions[start_idx:idx+1], test_pv_pred[start_idx:idx+1]))
        start_idx = idx + 1
    segments.append((times[start_idx:], predictions[start_idx:], test_pv_pred[start_idx:]))
    # 绘制多个子图
    num_segments = len(segments)
    fig, axes = plt.subplots(num_segments, 1, figsize=(14, 5 * num_segments), sharex=False)
    # 如果只有一个子图，将其转换为列表以便统一处理
    if num_segments == 1:
        axes = [axes]
    for i, (seg_times, seg_predictions, seg_test_pv_pred) in enumerate(segments):
        ax = axes[i]
        # 填充真实值曲线下的区域为灰色
        ax.fill_between(seg_times, seg_test_pv_pred, color='gray', alpha=0.5, label='True Values')
        # 绘制预测值曲线为绿色实线
        ax.plot(seg_times, seg_predictions, label='Predictions', color='green', linewidth=2)
        # 设置横坐标格式
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M'))  # 时间格式
        ax.xaxis.set_major_locator(mdates.AutoDateLocator())  # 自动选择时间间隔
        fig.autofmt_xdate()  # 自动旋转日期标签以避免重叠
        ax.set_title(f'Segment {i+1}: PV Predictions vs True Values Over Time', fontsize=14)
        ax.set_xlabel('Time', fontsize=12)
        ax.set_ylabel('Value', fontsize=12)
        ax.legend(fontsize=10)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Plot saved to {save_path}")



def calculate_rmse_BCD(metrics_list,mae_list, weather_list):
    # 提取 B (CLOUDY), C (SUNNY), D (OVERCAST) 的 RMSE 和 N
    rmse_A, N_A = metrics_list[3], weather_list[3]  # A (CLOUDY) 的 RMSE 和 N
    rmse_B, N_B = metrics_list[0], weather_list[0]  # B (CLOUDY) 的 RMSE 和 N
    rmse_C, N_C = metrics_list[1], weather_list[1]  # C (SUNNY) 的 RMSE 和 N
    rmse_D, N_D = metrics_list[2], weather_list[2]  # D (OVERCAST) 的 RMSE 和 N
    # 计算 BCD 的 MSE
    mse_B = rmse_B ** 2
    mse_C = rmse_C ** 2
    mse_D = rmse_D ** 2
    # 总样本数
    total_N = N_B + N_C + N_D
    # 加权平均 MSE
    mse_BCD = (N_B * mse_B + N_C * mse_C + N_D * mse_D) / total_N
    # 转换为 RMSE
    rmse_BCD = np.sqrt(mse_BCD)
    # 假设 metrics_list 中除了 RMSE 还包含 MAE，例如 metrics_list = [(rmse_B, mae_B), (rmse_C, mae_C), (rmse_D, mae_D)]
    mae_A = mae_list[3]
    mae_B = mae_list[0]  # B (CLOUDY) 的 MAE
    mae_C = mae_list[1]  # C (SUNNY) 的 MAE
    mae_D = mae_list[2]  # D (OVERCAST) 的 MAE
    # 加权平均 MAE
    mae_BCD = (N_B * mae_B + N_C * mae_C + N_D * mae_D ) / total_N
    # 输出结果
    print(f"MAE for BCD (CLOUDY, SUNNY, OVERCAST): {mae_BCD:.3f}")
    print(f"MAE for A: {mae_A:.3f}")
    # 输出结果
    print(f"RMSE for BCD (CLOUDY, SUNNY, OVERCAST): {rmse_BCD:.3f}")
    print(f"RMSE for A: {rmse_A:.3f}")
    return rmse_BCD, rmse_A